package main

import (
    "encoding/csv"
    "flag"
    "fmt"
    "log"
    "math/rand"
    "os"
    "path/filepath"
    "strconv"
    "strings"
    "time"
	"sort"
)

var (
    mongoURI       string
    databaseName   string
    collectionName string
    BaseTime       time.Time
    testCount      int
    queryMode      string
    allowDiskUse   bool
    sensorCount    int
    thresholds     []float64
	timeWindows     []time.Duration
	saveResults     bool
)

func parseArgs() {
    var baseTimeStr string
    var thresholdsStr string
	var windowsStr string

    flag.StringVar(&mongoURI, "uri", "mongodb://localhost:27017", "MongoDB 连接 URI")
    flag.StringVar(&databaseName, "db", "testdb", "数据库名称")
    flag.StringVar(&collectionName, "col", "testcol", "集合名称")
    flag.StringVar(&baseTimeStr, "base-time", "", "起始时间 (UTC)，格式: 2025-08-19T01:42:44Z")
    flag.IntVar(&testCount, "test-count", 20, "每个时间窗口的查询次数")
    flag.StringVar(&queryMode, "query-mode", "aggregate", "查询模式: count 或 aggregate")
    flag.BoolVar(&allowDiskUse, "allow-disk-use", false, "聚合时是否允许磁盘")
    flag.IntVar(&sensorCount, "sensor-count", 500, "可随机生成的传感器数量")
    flag.StringVar(&thresholdsStr, "thresholds", "1,2,5,10,20,50,100", "触发查询的磁盘阈值 (GB, 逗号分隔)")
	flag.StringVar(&windowsStr, "windows", "10s,1h", "查询时间窗口，逗号分隔，例如: 10s,1h,24h")
	flag.BoolVar(&saveResults, "save-results", false, "是否保存查询结果到 CSV 文件")

    flag.Parse()

    // 设置 BaseTime
    if baseTimeStr == "" {
        BaseTime = time.Now().UTC()
    } else {
        var err error
        BaseTime, err = time.Parse(time.RFC3339, baseTimeStr)
        if err != nil {
            log.Fatalf("解析 base-time 参数失败: %v", err)
        }
    }

    // 解析 thresholds
    for _, v := range strings.Split(thresholdsStr, ",") {
        f, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
        if err != nil {
            log.Fatalf("解析 thresholds 失败: %v", err)
        }
        thresholds = append(thresholds, f)
    }

	// 解析时间窗口
    windows := strings.Split(windowsStr, ",")
    for _, w := range windows {
        d, err := time.ParseDuration(strings.TrimSpace(w))
        if err != nil {
            log.Fatalf("无效的时间窗口 '%s'，请使用 10s、1h、24h 格式: %v", w, err)
        }
        timeWindows = append(timeWindows, d)
    }

    fmt.Println(">>> 参数配置")
    fmt.Printf("Mongo URI     : %s\n", mongoURI)
    fmt.Printf("Database      : %s\n", databaseName)
    fmt.Printf("Collection    : %s\n", collectionName)
    fmt.Printf("Base Time     : %s\n", BaseTime.Format(time.RFC3339))
    fmt.Printf("Thresholds    : %v GB\n", thresholds)
    fmt.Printf("Time Windows  : %v\n", timeWindows)
    fmt.Printf("Sensor Count  : %d\n", sensorCount)
    fmt.Println("====================================\n")
}

func randomSensorID() string {
    return fmt.Sprintf("sensor-%d", rand.Intn(sensorCount))
}

func createCSVFile() (*os.File, *csv.Writer) {
    os.MkdirAll("results", os.ModePerm)
    filename := filepath.Join("results", fmt.Sprintf("query_results_%d.csv", time.Now().Unix()))
    file, err := os.Create(filename)
    if err != nil {
        log.Fatalf("创建 CSV 文件失败: %v", err)
    }

    writer := csv.NewWriter(file)
    writer.Write([]string{"SensorID", "Timestamp", "Temperatures"})
    writer.Flush()
    return file, writer
}

func calculateStats(durations []float64) (avg, min, max, p95 float64) {
    if len(durations) == 0 {
        return 0, 0, 0, 0
    }

    // 排序
    sorted := make([]float64, len(durations))
    copy(sorted, durations)
    sort.Float64s(sorted)

    sum := 0.0
    for _, d := range sorted {
        sum += d
    }
    avg = sum / float64(len(sorted))
    min = sorted[0]
    max = sorted[len(sorted)-1]

    // P95
    idx := int(float64(len(sorted))*0.95) - 1
    if idx < 0 {
        idx = 0
    } else if idx >= len(sorted) {
        idx = len(sorted) - 1
    }
    p95 = sorted[idx]
    return
}